import copy
import json
import math
import struct

import numpy as np

from simple_ML_model import MY_NET_one_hiddenlayer_NN


def _data_norm(data):
    mean = np.mean(data)
    std_variance = math.sqrt(np.var(data))
    norm_data = (data - mean) / std_variance
    return norm_data

def data_norm(data):
    data_list = []
    for i in range(len(data)):
        norm_data_list = _data_norm(data[i]).tolist()  #每次对1 * 784大小的nparray做标准化（normalization）
        data_list.append(norm_data_list)
    norm_data = np.asarray(data_list)
        
    return norm_data
    
#改成自己的路径
with open('C:\\Users\\林俊锞\\Desktop\\mnist\\binary_mnist\\t10k-labels.idx1-ubyte', mode = 'rb') as test_lable_file:
    
    magic , num_of_test_lables_item = struct.unpack('>II', test_lable_file.read(8)) #前面8个字节不是真正的数据
    test_y_total_np = np.fromfile(test_lable_file, dtype = np.uint8) 
    test_y_np = test_y_total_np[0 : 3000]
    
with open('C:\\Users\\林俊锞\\Desktop\\mnist\\binary_mnist\\t10k-images.idx3-ubyte', mode = 'rb') as test_image_file:
    magic , num_of_test_images_item, rows, columns = struct.unpack('>IIII', test_image_file.read(16)) #前16个字节不是真正的数据
    test_x_total_np = np.fromfile(test_image_file, dtype = np.uint8).reshape(num_of_test_images_item, 28 * 28)
    test_x_np = data_norm(test_x_total_np[0 : 3000])

with open('C:\\Users\\林俊锞\\Desktop\\mnist\\binary_mnist\\train-images.idx3-ubyte', mode = 'rb') as train_image_file:
    magic, num_of_train_images_item, _, _ = struct.unpack('>IIII', train_image_file.read(16))
    train_x_total_np = np.fromfile(train_image_file, dtype = np.uint8).reshape(num_of_train_images_item, 28 * 28)
    train_x_np = data_norm(train_x_total_np[0 : 30000])

with open('C:\\Users\\林俊锞\\Desktop\\mnist\\binary_mnist\\train-labels.idx1-ubyte', mode = 'rb') as train_label_file:
    magic, num_of_train_lables_item = struct.unpack('>II', train_label_file.read(8))
    train_y_total_np = np.fromfile(train_label_file, dtype = np.uint8)
    train_y_np = train_y_total_np[0 : 30000]

if __name__ == '__main__':
    epoch = 10
    lr = 0.05
    batch_size = 10
    w1_shape = (128, 28 * 28)
    w2_shape = (10, 128)
    net = MY_NET_one_hiddenlayer_NN(lr, batch_size, w1_shape, w2_shape)
    for i in range(epoch):
        print('epoch ', i)
        train_loss = net.train(train_x = train_x_np, train_y = train_y_np)
        acc = net.test(test_x = test_x_np, test_y = test_y_np)
        print(train_loss, acc)
